#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional, Tuple, Literal
import asyncio

from app.core.logging import get_logger, LoggerMixin
from app.core.exceptions import NotFoundError, DatabaseError, EmbeddingError, VectorDBError
from app.repositories.training_data import TrainingDataRepository
from app.repositories.vector import VectorRepository
from app.services.embedding import EmbeddingService
from app.models.database import TrainingData

# 创建日志记录器
logger = get_logger("training_data_service")


class TrainingDataService(LoggerMixin):
    """训练数据服务
    
    管理训练数据相关业务逻辑
    """
    
    def __init__(
        self,
        training_data_repo: TrainingDataRepository,
        vector_repo: VectorRepository,
        embedding_service: EmbeddingService
    ):
        """初始化训练数据服务
        
        Args:
            training_data_repo: 训练数据仓库
            vector_repo: 向量仓库
            embedding_service: 嵌入服务
        """
        self.training_data_repo = training_data_repo
        self.vector_repo = vector_repo
        self.embedding_service = embedding_service
    
    async def _get_vector_data_type(self, data_type: str) -> str:
        """获取向量数据类型
        
        Args:
            data_type: 数据类型
            
        Returns:
            str: 向量数据类型
        """
        # 现在所有类型都使用同一个集合
        return "train"
    
    async def _get_text_for_embedding(self, data: TrainingData) -> str:
        """获取用于嵌入的文本
        
        Args:
            data: 训练数据
            
        Returns:
            str: 用于嵌入的文本
        """
        if data.type == "ddl":
            return data.content
        elif data.type == "question_sql":
            return data.question or ""
        elif data.type == "description":
            return data.content
        else:
            return data.content
            
    async def _get_sql_for_embedding(self, data: TrainingData) -> Optional[str]:
        """获取用于嵌入的SQL
        
        Args:
            data: 训练数据
            
        Returns:
            Optional[str]: 用于嵌入的SQL，如果不是question_sql类型则返回None
        """
        if data.type == "question_sql" and data.sql:
            return data.sql
        return None
    
    async def add_training_data(
        self,
        project_id: str,
        data_type: Literal["ddl", "question_sql", "description"],
        content: str,
        question: Optional[str] = None,
        sql: Optional[str] = None
    ) -> Dict[str, Any]:
        """添加训练数据
        
        Args:
            project_id: 项目ID
            data_type: 数据类型
            content: 内容
            question: 问题
            sql: SQL查询
            
        Returns:
            Dict[str, Any]: 添加的训练数据字典
            
        Raises:
            DatabaseError: 当添加失败时
            EmbeddingError: 当生成嵌入失败时
            VectorDBError: 当存储向量失败时
        """
        try:
            # 创建训练数据
            training_data_dict = await self.training_data_repo.create(
                project_id=project_id,
                data_type=data_type,
                content=content,
                question=question,
                sql=sql
            )
            
            # 确保项目的向量集合存在
            await self.vector_repo.ensure_collections(project_id)
            
            # 获取用于嵌入的文本
            text = ""
            if data_type == "ddl":
                text = content
            elif data_type == "question_sql":
                text = question or ""
            elif data_type == "description":
                text = content
            
            # 生成嵌入
            vector = await self.embedding_service.embed_text(text)
            
            # 获取向量数据类型
            vector_data_type = await self._get_vector_data_type(data_type)
            
            # 存储向量
            await self.vector_repo.insert_vector(
                project_id=project_id,
                data_type=vector_data_type,
                content_id=training_data_dict["id"],
                vector=vector,
                content_type=data_type  # 传递内容类型
            )
            
            # 对于question_sql类型，还需要存储SQL的向量表示
            if data_type == "question_sql" and sql:
                # 生成SQL嵌入
                sql_vector = await self.embedding_service.embed_text(sql)
                
                # 存储SQL向量，使用相同的集合但不同的content_type
                await self.vector_repo.insert_vector(
                    project_id=project_id,
                    data_type="train",  # 使用统一的data_type
                    content_id=training_data_dict["id"],
                    vector=sql_vector,
                    content_type="sql"  # 标记为SQL类型
                )
                self.logger.info(f"已添加SQL向量: {training_data_dict['id']}")
            
            self.logger.info(f"已添加训练数据: {training_data_dict['id']}")
            return training_data_dict
            
        except Exception as e:
            self.logger.error(f"添加训练数据失败: {str(e)}")
            raise
    
    async def get_training_data(self, data_id: int) -> Dict[str, Any]:
        """获取训练数据
        
        Args:
            data_id: 训练数据ID
            
        Returns:
            Dict[str, Any]: 训练数据字典
            
        Raises:
            NotFoundError: 当训练数据不存在时
        """
        try:
            training_data_dict = await self.training_data_repo.get_by_id(data_id)
            return training_data_dict
        except Exception as e:
            self.logger.error(f"获取训练数据失败: {str(e)}")
            raise
    
    async def delete_training_data(self, data_id: int) -> None:
        """删除训练数据
        
        Args:
            data_id: 训练数据ID
            
        Raises:
            NotFoundError: 当训练数据不存在时
            DatabaseError: 当删除失败时
            VectorDBError: 当删除向量失败时
        """
        try:
            # 获取训练数据
            training_data_dict = await self.training_data_repo.get_by_id(data_id)
            
            # 获取向量数据类型
            data_type = training_data_dict["type"] if isinstance(training_data_dict, dict) else training_data_dict.type
            vector_data_type = await self._get_vector_data_type(data_type)
            
            # 获取项目ID
            project_id = training_data_dict["project_id"] if isinstance(training_data_dict, dict) else training_data_dict.project_id
            
            # 删除向量
            try:
                # 删除主要向量
                await self.vector_repo.delete_vector(
                    project_id=project_id,
                    data_type=vector_data_type,
                    content_id=data_id
                )
                
                # 如果是question_sql类型，还需要删除SQL向量
                if data_type == "question_sql":
                    try:
                        await self.vector_repo.delete_vector(
                            project_id=project_id,
                            data_type="ddl",  # SQL向量存储在DDL集合中
                            content_id=data_id
                        )
                    except Exception as e:
                        self.logger.error(f"删除SQL向量失败，但将继续删除训练数据: {str(e)}")
            except Exception as e:
                self.logger.error(f"删除向量失败，但将继续删除训练数据: {str(e)}")
            
            # 删除训练数据
            await self.training_data_repo.delete(data_id)
            
            self.logger.info(f"已删除训练数据: {data_id}")
            
        except Exception as e:
            self.logger.error(f"删除训练数据失败: {str(e)}")
            raise
    
    async def list_training_data(
        self,
        project_id: str,
        data_type: Optional[str] = None,
        limit: int = 10,
        offset: int = 0
    ) -> Tuple[List[Dict[str, Any]], int]:
        """列出训练数据
        
        Args:
            project_id: 项目ID
            data_type: 数据类型
            limit: 限制数量
            offset: 偏移量
            
        Returns:
            Tuple[List[Dict[str, Any]], int]: 训练数据字典列表和总数
            
        Raises:
            DatabaseError: 当查询失败时
        """
        try:
            training_data_list, total = await self.training_data_repo.list_by_project(
                project_id=project_id,
                data_type=data_type,
                limit=limit,
                offset=offset
            )
            return training_data_list, total
        except Exception as e:
            self.logger.error(f"列出训练数据失败: {str(e)}")
            raise
    
    async def get_similar_questions(
        self,
        project_id: str,
        question: str,
        top_k: int = 5
    ) -> List[TrainingData]:
        """获取相似问题
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相似问题列表
            
        Raises:
            EmbeddingError: 当生成嵌入失败时
            VectorDBError: 当搜索向量失败时
            DatabaseError: 当查询失败时
        """
        try:
            # 生成问题嵌入
            vector = await self.embedding_service.embed_query(question)
            
            # 搜索相似向量
            similar_items = await self.vector_repo.search_similar(
                project_id=project_id,
                data_type="questions",
                vector=vector,
                top_k=top_k
            )
            
            if not similar_items:
                return []
            
            # 获取内容ID列表
            content_ids = [item["content_id"] for item in similar_items]
            
            # 获取问题-SQL数据
            question_sql_list = await self.training_data_repo.get_question_sql_by_ids(content_ids)
            
            return question_sql_list
            
        except Exception as e:
            self.logger.error(f"获取相似问题失败: {str(e)}")
            raise
    
    async def get_relevant_ddl(
        self,
        project_id: str,
        question: str,
        top_k: int = 3
    ) -> List[TrainingData]:
        """获取相关DDL
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相关DDL列表
            
        Raises:
            DatabaseError: 当查询失败时
        """
        try:
            # 获取所有DDL
            ddl_list = await self.training_data_repo.get_ddl_by_project(project_id)
            
            if not ddl_list:
                return []
            
            # 如果DDL数量小于等于top_k，直接返回所有DDL
            if len(ddl_list) <= top_k:
                return ddl_list
            
            # 生成问题嵌入
            question_vector = await self.embedding_service.embed_query(question)
            
            # 生成所有DDL的嵌入
            ddl_texts = [ddl.content for ddl in ddl_list]
            ddl_vectors = await self.embedding_service.embed_texts(ddl_texts)
            
            # 计算相似度
            similarities = []
            for i, ddl_vector in enumerate(ddl_vectors):
                similarity = self.embedding_service.calculate_similarity(question_vector, ddl_vector)
                similarities.append((i, similarity))
            
            # 按相似度排序
            similarities.sort(key=lambda x: x[1], reverse=True)
            
            # 获取top_k个相关DDL
            relevant_ddl = [ddl_list[i] for i, _ in similarities[:top_k]]
            
            return relevant_ddl
            
        except Exception as e:
            self.logger.error(f"获取相关DDL失败: {str(e)}")
            raise
    
    async def get_relevant_descriptions(
        self,
        project_id: str,
        question: str,
        top_k: int = 3
    ) -> List[TrainingData]:
        """获取相关描述
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相关描述列表
            
        Raises:
            EmbeddingError: 当生成嵌入失败时
            VectorDBError: 当搜索向量失败时
            DatabaseError: 当查询失败时
        """
        try:
            # 生成问题嵌入
            vector = await self.embedding_service.embed_query(question)
            
            # 搜索相似向量
            similar_items = await self.vector_repo.search_similar(
                project_id=project_id,
                data_type="descriptions",
                vector=vector,
                top_k=top_k
            )
            
            if not similar_items:
                return []
            
            # 获取内容ID列表
            content_ids = [item["content_id"] for item in similar_items]
            
            # 获取描述数据
            description_list = await self.training_data_repo.get_descriptions_by_ids(content_ids)
            
            return description_list
            
        except Exception as e:
            self.logger.error(f"获取相关描述失败: {str(e)}")
            raise